package com.example.scg.route;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import lombok.extern.slf4j.Slf4j;
import org.springframework.cloud.gateway.event.RefreshRoutesEvent;
import org.springframework.cloud.gateway.route.RouteDefinition;
import org.springframework.cloud.gateway.route.RouteDefinitionRepository;
import org.springframework.context.ApplicationEventPublisher;
import org.springframework.context.ApplicationEventPublisherAware;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

/**
 * @author Aaron
 */
@Slf4j
@Component
public class MongoRouteDefinitionRepository
    implements RouteDefinitionRepository, ApplicationEventPublisherAware {

  private ApplicationEventPublisher eventPublisher;

  private Map<String, RouteDefinition> cache = new ConcurrentHashMap<>();

  private final RouteRepositoryOperations repositoryOperation;

  public MongoRouteDefinitionRepository(RouteRepositoryOperations repositoryOperation) {
    this.repositoryOperation = repositoryOperation;
  }

  @Override
  public Flux<RouteDefinition> getRouteDefinitions() {
    return Flux.fromIterable(cache.values());
  }

  @Override
  public Mono<Void> save(Mono<RouteDefinition> route) {
    return route.flatMap(
        r -> repositoryOperation.save(MongoRouteDefinition.from(r))
            .log()
            .doOnNext(this::addCache)
            .then(Mono.empty())
    );
  }

  @Override
  public Mono<Void> delete(Mono<String> routeId) {
    return repositoryOperation.findById(routeId)
        .log()
        .map(RouteDefinition::getId)
        .doOnNext(this::removeCache)
        .flatMap(repositoryOperation::deleteById);
  }

  @Override
  public void setApplicationEventPublisher(ApplicationEventPublisher eventPublisher) {
    this.eventPublisher = eventPublisher;
  }

  /**
   * 将指定路由加入到缓存中。
   * <p>
   * 为了能实时加载路由，可以通过MongoDB的ChangeStream,监听到数据变化后调用此方法
   */
  public void addCache(RouteDefinition route) {
    this.cache.putIfAbsent(route.getId(), route);
    this.publishEvent();
  }

  /**
   * 将指定路由从缓存中删除。
   * <p>
   * 为了能实时加载路由，可以通过MongoDB的ChangeStream,监听到数据变化后调用此方法
   */
  public void removeCache(String routeId) {
    if (this.cache.remove(routeId) != null) {
      this.publishEvent();
    }
  }

  void publishEvent() {
    eventPublisher.publishEvent(new RefreshRoutesEvent(this));
  }

  RouteRepositoryOperations getRepositoryOperation() {
    return repositoryOperation;
  }


  Map<String, RouteDefinition> getCache() {
    return cache;
  }

  void setCache(
      Map<String, RouteDefinition> cache) {
    this.cache = cache;
  }

}
